# 库函数调用
import torch
from torch import nn
import torch.nn.functional as F

import numpy as np

# from measure import Four_three
#https://blog.csdn.net/wangshuqian1314/article/details/139632509?utm_medium=distribute.pc_relevant.none-task-blog-2~default~baidujs_baidulandingword~default-0-139632509-blog-133668315.235^v43^pc_blog_bottom_relevance_base7&spm=1001.2101.3001.4242.1&utm_relevant_index=3

# 三维卷积块定义
class DoubleConv(nn.Module):
    """(Conv3D -> IN -> ReLU) * 2"""

    def __init__(self, in_channels, out_channels, num_groups = 8):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size = 3, stride = 1, padding = 1,bias=True),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size = 3, stride = 1, padding = 1,bias=True),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

# 下采样模块定义
class Down(nn.Module):

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.MaxPool3d(2,2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.encoder(x)

# 上采样模块定义
class Up(nn.Module):

    def __init__(self, in_channels, out_channels, trilinear = True):
        super().__init__()

        if trilinear:
            self.up = nn.Upsample(scale_factor = 2)
        else:
            self.up = nn.ConvTranspose3d(in_channels // 2, in_channels // 2, kernel_size = 2, stride = 2)

        self.conv = DoubleConv(in_channels, out_channels)
        self.downc = nn.Conv3d(in_channels, out_channels, kernel_size = 3, stride=1, padding=1, bias=True)
        self.downr = nn.ReLU(inplace=True)

    def forward(self, x1, x2):
        x1 = self.up(x1)

        diffZ = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        diffX = x2.size()[4] - x1.size()[4]
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2, diffZ // 2, diffZ - diffZ // 2])

        x1 = self.downr(self.downc(x1))

        x = torch.cat([x2, x1], dim = 1)
        return self.conv(x)

# 输出卷积层定义
class Out(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size = 3, stride=1, padding=1)

    def forward(self, x):
        return self.conv(x)

# 3D-UNet模型定义
class UNET3D(nn.Module):
    def __init__(self, in_channels=3, out_channels=1,n_channels=64):
        super().__init__()
        self.in_channels = in_channels
        self.n_channels = n_channels

        self.conv = DoubleConv(in_channels, n_channels)
        self.enc1 = Down(n_channels, 2 * n_channels)
        self.enc2 = Down(2 * n_channels, 4 * n_channels)
        self.enc3 = Down(4 * n_channels, 8 * n_channels)
        self.enc4 = Down(8 * n_channels, 16 * n_channels)

        self.dec1 = Up(16 * n_channels, 8 * n_channels)
        self.dec2 = Up(8 * n_channels, 4 * n_channels)
        self.dec3 = Up(4 * n_channels, 2*n_channels)
        self.dec4 = Up(2 * n_channels, n_channels)
        self.out = Out(n_channels, out_channels) #(1,4,128,128,n)


    def forward(self, x):
        # print('size of x:', x.shape)
        x1 = self.conv(x)
        # print('size of x1:', x1.shape)
        x2 = self.enc1(x1)
        # print('size of x2:', x2.shape)
        x3 = self.enc2(x2)
        # print('size of x3:', x3.shape)
        x4 = self.enc3(x3)
        # print('size of x4:', x4.shape)
        x5 = self.enc4(x4)
        # print('size of x5:', x5.shape)

        mask = self.dec1(x5, x4)
        # print('size of mask:', mask.shape)
        mask = self.dec2(mask, x3)
        # print('size of mask:', mask.shape)
        mask = self.dec3(mask, x2)
        # print('size of mask:', mask.shape)
        mask = self.dec4(mask, x1)
        # print('size of mask:', mask.shape)
        mask = self.out(mask)
        # print('size of mask:', mask.shape)

        return mask

# 测试代码
if __name__ == '__main__':
    input_channels = 4
    output_channels = 1
    x = torch.ones([16, 4, 16, 16, 16])
    # 假设这里是随机生成的目标标签，实际应用中需要替换为真实标签
    y_true = torch.randint(0, 2, [16, 1, 16, 16, 16]).float()
    model = UNET3D(input_channels, output_channels)
    print('model initialization finished!')

    # 定义损失函数和优化器
    criterion = torch.nn.BCEWithLogitsLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    # 训练参数
    num_epochs = 10

    for epoch in range(num_epochs):
        # 前向传播
        f = model(x)
        loss = criterion(f, y_true)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 打印训练信息
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

    # 训练完成后，再次计算输出的均值和标准差
    f = model(x)
    mean = torch.mean(f)
    std = torch.std(f)
    print(f"Mean of the model output: {mean}")
    print(f"Standard deviation of the model output: {std}")